"""
新闻情绪分析策略
策略逻辑：
1. 实时获取财经新闻
2. 使用NLP分析情绪得分
3. 根据情绪变化生成交易信号
4. 结合技术指标确认信号
"""

import pandas as pd
from transformers import pipeline
from datetime import datetime, timedelta

class SentimentTradingStrategy:
    def __init__(self, ticker, lookback_days=30):
        self.ticker = ticker
        self.lookback = lookback_days
        self.sentiment_analyzer = pipeline(
            "sentiment-analysis",
            model="finiteautomata/bertweet-base-sentiment-analysis"
        )
        self.sentiment_history = pd.DataFrame(
            columns=['date', 'sentiment', 'score']
        )
        
    def analyze_news(self, news_text):
        """分析新闻情绪"""
        result = self.sentiment_analyzer(news_text[:512])[0]  # 限制文本长度
        sentiment_map = {'POS': 1, 'NEU': 0, 'NEG': -1}
        return {
            'sentiment': sentiment_map[result['label']],
            'score': result['score']
        }
    
    def update_sentiment(self, news_items):
        """更新情绪历史数据"""
        for item in news_items:
            analysis = self.analyze_news(item['content'])
            self.sentiment_history = self.sentiment_history.append({
                'date': datetime.now(),
                'sentiment': analysis['sentiment'],
                'score': analysis['score']
            }, ignore_index=True)
        
        # 保留最近30天数据
        cutoff = datetime.now() - timedelta(days=self.lookback)
        self.sentiment_history = self.sentiment_history[
            self.sentiment_history['date'] > cutoff
        ]
    
    def calculate_sentiment_index(self):
        """计算情绪指数(7日加权平均)"""
        if len(self.sentiment_history) == 0:
            return 0
            
        recent = self.sentiment_history.tail(7)
        return (recent['sentiment'] * recent['score']).mean()
    
    def generate_signal(self, current_price, ma20):
        """生成交易信号"""
        sentiment_index = self.calculate_sentiment_index()
        
        # 多头信号: 情绪积极且价格在20日均线上方
        if sentiment_index > 0.5 and current_price > ma20:
            return 1
        # 空头信号: 情绪消极且价格在20日均线下方
        elif sentiment_index < -0.5 and current_price < ma20:
            return -1
        else:
            return 0

if __name__ == '__main__':
    strategy = SentimentTradingStrategy('AAPL')
    
    # 模拟新闻数据
    test_news = [
        {'content': 'Apple reports record quarterly earnings'},
        {'content': 'New iPhone sales exceed expectations'},
        {'content': 'Supply chain issues may affect production'}
    ]
    
    # 更新情绪数据
    strategy.update_sentiment(test_news)
    
    # 生成信号
    signal = strategy.generate_signal(150, 145)
    print(f"交易信号: {signal}")